{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook can be executed in a notebook hosted in KubeFlow. \n",
    "You can find instructions on how to deploy a KubeFlow cluster and how to access the the KubeFlow UI and the hosted notebooks here: https://www.kubeflow.org/docs/pipelines/pipelines-quickstart/\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please install KubeFlow Pipelines SDK using the following comand:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install 'https://storage.googleapis.com/ml-pipeline/release/0.1.9/kfp.tar.gz'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Energy Price Forecasting Pipeline\n",
    "\n",
    "This notebook generates a KubeFlow pipeline that runs the solution end to end.\n",
    "For more information on KubeFlow pipelines and how to run them in GCP please visit https://github.com/kubeflow/pipelines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import kfp\n",
    "from kfp import compiler\n",
    "import kfp.dsl as dsl\n",
    "import kfp.gcp as gcp\n",
    "import kfp.notebook\n",
    "\n",
    "#Please modify the following values to match your GCP bucket, project, and docker image name.\n",
    "OUTPUT_DIR = 'gs://pipelinestest/out'\n",
    "PROJECT_NAME = 'energy-forecasting'\n",
    "EF_IMAGE='gcr.io/%s/energy:dev' % PROJECT_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create base image\n",
    "\n",
    "This image takes the `tensorflow/tensorflow:1.10.0-py3` as a starting point and installs python libraries and applications that are required by some components in the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%docker {EF_IMAGE} {OUTPUT_DIR}\n",
    "FROM tensorflow/tensorflow:1.10.0-py3\n",
    "RUN apt-get update \n",
    "RUN apt-get install -y git\n",
    "RUN pip3 install --upgrade google-api-python-client\n",
    "RUN pip3 install --upgrade pyarrow\n",
    "RUN pip3 install --upgrade google-cloud-bigquery\n",
    "RUN pip3 install --upgrade google-cloud-storage\n",
    "RUN pip3 install --upgrade gitpython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Components\n",
    "\n",
    "Each cell defines the logic of different components that will be used in the pipeline and produces a `.yaml` file for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def copy_table(\n",
    "    dataset: str) -> str:\n",
    "    \"\"\"Retrieves raw data from competition website.\n",
    "    Retrieves raw data from the competition site and saves it in BigQuery.\n",
    "    Args:\n",
    "        dataset: String specifying the dataset in BigQuery to save the data in.\n",
    "    Returns:\n",
    "        String specifying if the component finished succesfully.\n",
    "    \"\"\"\n",
    "    from google.cloud import bigquery\n",
    "    import requests\n",
    "    import pandas as pd\n",
    "    from io import StringIO\n",
    "    from io import BytesIO\n",
    "    import zipfile\n",
    "    \n",
    "    bq_client = bigquery.Client()\n",
    "    price_data = pd.read_csv(\n",
    "        StringIO(requests.get(\n",
    "            'http://complatt.smartwatt.net/assets/files/historicalRealData/RealMarketPriceDataPT.csv').text),\n",
    "        sep=';'\n",
    "    )\n",
    "    price_data.columns = ['date_utc', 'price']\n",
    "    bq_client.load_table_from_dataframe(\n",
    "        price_data,\n",
    "        bq_client.dataset(dataset).table(\n",
    "            'MarketPricePT')).result()\n",
    "    \n",
    "    weather_zip = zipfile.ZipFile(\n",
    "        BytesIO(requests.get(\n",
    "            'http://complatt.smartwatt.net/assets/files/weatherHistoricalData/WeatherHistoricalData.zip').content))\n",
    "    weather_data = pd.read_csv(\n",
    "        weather_zip.open(\n",
    "            'WeatherHistoricalData/historical_weather.csv'))\n",
    "    bq_client.load_table_from_dataframe(\n",
    "        weather_data,\n",
    "        bq_client.dataset(dataset).table(\n",
    "            'historical_weather')).result()\n",
    "    \n",
    "    return('success')\n",
    "    \n",
    "    \n",
    "compiler.build_python_component(\n",
    "    component_func = copy_table,\n",
    "    staging_gcs_path = OUTPUT_DIR,\n",
    "    base_image=EF_IMAGE,\n",
    "    target_component_file='copy-table.component.yaml',\n",
    "    target_image = 'gcr.io/' + PROJECT_NAME + '/component-copy-table:latest')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def export_table(\n",
    "    inp: str,\n",
    "    table: str,\n",
    "    file: str) -> str:\n",
    "    \"\"\"Exports table to csv.\n",
    "    Exports BigQuery table into CSV file.\n",
    "    Args:\n",
    "        inp: String containing the output from previous component.\n",
    "        table: String specifying the origin BigQuery table.\n",
    "        file: String specifying the path and name for the csv file.\n",
    "    Returns:\n",
    "        String specifying if the component finished succesfully.\n",
    "    \"\"\"\n",
    "    from google.cloud import bigquery\n",
    "    bq_client = bigquery.Client()\n",
    "    bq_client.extract_table(\n",
    "        table,\n",
    "        file).result()\n",
    "    return('success')\n",
    "    \n",
    "compiler.build_python_component(\n",
    "    component_func = export_table,\n",
    "    staging_gcs_path = OUTPUT_DIR,\n",
    "    base_image=EF_IMAGE,\n",
    "    target_component_file='export-table.component.yaml',\n",
    "    target_image = 'gcr.io/' + PROJECT_NAME + '/component-export-table:latest')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_git_python_script(\n",
    "    inp: str,\n",
    "    code_repo: str,\n",
    "    code_folder: str,\n",
    "    script: str,\n",
    "    script_args: str) -> str:\n",
    "    \"\"\"Runs Python script from git repository.\n",
    "\n",
    "    Args:\n",
    "        inp: String containing the output from previous component.\n",
    "        code_repo: String specifying the url to the git repository.\n",
    "        code_folder: String specifying the folder for the script.\n",
    "        script: String specifying the name of the script.\n",
    "        script_args: String specifying the arguments for the script.\n",
    "    Returns:\n",
    "        String specifying if the component finished succesfully.\n",
    "    \"\"\"\n",
    "    import os\n",
    "    import git\n",
    "    git.Git('').clone(code_repo)\n",
    "    os.chdir(code_folder)\n",
    "    output = os.system(' '.join([\n",
    "        'python -m',\n",
    "        script,\n",
    "        script_args]))\n",
    "    if output == 0:\n",
    "        return('success')\n",
    "    raise Exception('Script failed. The exit status was: {}'.format(output))\n",
    "    \n",
    "compiler.build_python_component(\n",
    "    component_func = run_git_python_script,\n",
    "    staging_gcs_path = OUTPUT_DIR,\n",
    "    base_image=EF_IMAGE,\n",
    "    target_component_file='run-git-python-script.component.yaml',\n",
    "    target_image = 'gcr.io/' + PROJECT_NAME + '/component-run-git-python-script:latest')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_git_cmle_model(\n",
    "    tr_inp: str,\n",
    "    va_inp: str,\n",
    "    code_repo: str,\n",
    "    code_folder: str,\n",
    "    project: str,\n",
    "    bucket: str,\n",
    "    package_folder: str,\n",
    "    cmle_folder: str,\n",
    "    scale_tier: str,\n",
    "    python_module: str,\n",
    "    region: str,\n",
    "    runtime_version: str,\n",
    "    cmle_args: str) -> str:\n",
    "    \"\"\"Executes CMLE training job.\n",
    "    Retrieves python file from git repo and launches training job in CMLE.\n",
    "    Args:\n",
    "        tr_inp: String containing the source for the training data.\n",
    "        va_inp: String containing the source for the validation data.\n",
    "        code_repo: String specifying the url to the git repository.\n",
    "        code_folder: String specifying the folder for the job code.\n",
    "        project: String specifying the GCP project where job will run.\n",
    "        bucket: String specifying the GCS bucket where to save the job's outputs.\n",
    "        package_folder: String specifying the python package to run for the job.\n",
    "        cmle_folder: String specifying the folder in GCS where to save outputs.\n",
    "        scale_tier: String specifying compute resources to use for training job.\n",
    "        python_module: String specifying the python module to run for the job.\n",
    "        region: String specifying the GCP region in which to run the job.\n",
    "        runtime_version: String specifying the CMLE version to use for the job.\n",
    "        script_args: String specifying the arguments for the CMLE job.\n",
    "    Returns:\n",
    "        String containing output from running the training job in CMLE.\n",
    "    \"\"\"\n",
    "    import os\n",
    "    import git\n",
    "    import tarfile\n",
    "    import datetime\n",
    "    from google.cloud import storage\n",
    "    from googleapiclient import discovery\n",
    "    jobId = 'train' + datetime.datetime.today().strftime('%Y%m%d%H%M%S')\n",
    "    git.Git('').clone(code_repo)\n",
    "    with tarfile.open('code.tar.gz', 'w:gz') as tar:\n",
    "        tar.add(\n",
    "            code_folder,\n",
    "            arcname=os.path.basename(code_folder))\n",
    "    gcs_client = storage.Client()\n",
    "    gcs_bucket = gcs_client.get_bucket(bucket)\n",
    "    blob = gcs_bucket.blob(package_folder + jobId + '.tar.gz')\n",
    "    blob.upload_from_filename('code.tar.gz')\n",
    "    training_inputs = {\n",
    "        'scaleTier': scale_tier,\n",
    "        'pythonModule': python_module,\n",
    "        'args': cmle_args.split(' '),\n",
    "        'region': region,\n",
    "        'packageUris': [\n",
    "            'gs://'+ bucket + '/' + package_folder + jobId + '.tar.gz'],\n",
    "        'jobDir': 'gs://'+ bucket + '/' + cmle_folder + jobId,\n",
    "        'runtimeVersion': runtime_version}\n",
    "    job_spec = {\n",
    "        'jobId': jobId,\n",
    "        'trainingInput': training_inputs}\n",
    "    cloudml = discovery.build('ml', 'v1')\n",
    "    project_id = 'projects/{}'.format(project)\n",
    "    request = cloudml.projects().jobs().create(\n",
    "        body=job_spec,\n",
    "        parent=project_id)\n",
    "    return(str(request.execute()))\n",
    "\n",
    "    \n",
    "compiler.build_python_component(\n",
    "    component_func = train_git_cmle_model,\n",
    "    staging_gcs_path = OUTPUT_DIR,\n",
    "    base_image=EF_IMAGE,\n",
    "    target_component_file='train-git-cmle-model.component.yaml',\n",
    "    target_image = 'gcr.io/' + PROJECT_NAME + '/component-train-git-cmle-model:latest')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create pipeline\n",
    "\n",
    "The following code loads all components needed for the pipeline. Specifies dependencies between components. Defines arguments and defaults for the pipeline and saves the pipeline into a `.tar.gz` file that can be loaded into KubeFlow pipelines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dsl.pipeline(\n",
    "    name='Energy Price Forecasting',\n",
    "    description='Energy Price Forecasting')\n",
    "def basic_bq_pipeline(\n",
    "    project = dsl.PipelineParam(\n",
    "        'project',\n",
    "        value='energy-forecasting'),\n",
    "    dataset = dsl.PipelineParam(\n",
    "        'dataset',\n",
    "        value='Energy'),\n",
    "    bucket = dsl.PipelineParam(\n",
    "        'bucket',\n",
    "        value='energyforecast'),\n",
    "    code_repo = dsl.PipelineParam(\n",
    "        'code-repo',\n",
    "        value='https://github.com/GoogleCloudPlatform/professional-services.git'),\n",
    "    code_folder = dsl.PipelineParam(\n",
    "        'code-folder',\n",
    "        value='professional-services/examples/cloudml-energy-price-forecasting'),\n",
    "    data_prep_script = dsl.PipelineParam(\n",
    "        'data-prep-script',\n",
    "        value='data_preparation.data_prep'),\n",
    "    data_prep_args = dsl.PipelineParam(\n",
    "        'data-prep-args',\n",
    "        value=' '.join([\n",
    "            '--dataset=Energy',\n",
    "            '--train_table=MLDataTrain',\n",
    "            '--valid_table=MLDataValid',\n",
    "            '--test_table=MLDataTest',\n",
    "            '--prepare_data_file=data_preparation/prepare_data.sql',\n",
    "            '--weather_mean_std_file=data_preparation/weather_mean_std.sql',\n",
    "            '--train_from_date=\"2015-01-05 00:00:00\"',\n",
    "            '--train_to_date=\"2015-10-04 23:01:00\"',\n",
    "            '--valid_from_date=\"2015-10-05 00:00:00\"',\n",
    "            '--valid_to_date=\"2015-10-11 23:01:00\"',\n",
    "            '--test_from_date=\"2015-10-12 00:00:00\"',\n",
    "            '--test_to_date=\"2015-10-18 23:01:00\"',\n",
    "            '--price_scaling=0.01',\n",
    "            '--mean_path=gs://energyforecast/data/pickle/mean.pkl',\n",
    "            '--std_path=gs://energyforecast/data/pickle/std.pkl'])),\n",
    "    package_folder = dsl.PipelineParam(\n",
    "        'package-folder',\n",
    "        value='package/'),\n",
    "    cmle_folder = dsl.PipelineParam(\n",
    "        'cmle-folder',\n",
    "        value='cmle/'),\n",
    "    cmle_args = dsl.PipelineParam(\n",
    "        'cmle-args',\n",
    "        value=' '.join([\n",
    "            '--training_path', 'gs://energyforecast/data/csv/MLDataTrain.csv',\n",
    "            '--validation_path', 'gs://energyforecast/data/csv/MLDataValid.csv',\n",
    "            '--mean_path', 'gs://energyforecast/data/pickle/mean.pkl',\n",
    "            '--std_path', 'gs://energyforecast/data/pickle/std.pkl',\n",
    "            '--dropout' , '0.2',\n",
    "            '--hour_embedding', '20',\n",
    "            '--day_embedding', '10',\n",
    "            '--first_layer_size', '100',\n",
    "            '--number_layers', '3',\n",
    "            '--layer_reduction_fraction', '0.5',\n",
    "            '--learning_rate', '0.01',\n",
    "            '--batch_size', '64',\n",
    "            '--eval_batch_size', '168',\n",
    "            '--max_steps', '5000'])),\n",
    "    scale_tier = dsl.PipelineParam(\n",
    "        'scale-tier',\n",
    "        value='BASIC'),\n",
    "    python_module = dsl.PipelineParam(\n",
    "        'python-module',\n",
    "        value='trainer.task'),\n",
    "    region = dsl.PipelineParam(\n",
    "        'region',\n",
    "        value='us-central1'),\n",
    "    runtime_version = dsl.PipelineParam(\n",
    "        'runtime-version',\n",
    "        value='1.10'),\n",
    "    train_table = dsl.PipelineParam(\n",
    "        'train-table',\n",
    "        value='Energy.MLDataTrain'),\n",
    "    valid_table = dsl.PipelineParam(\n",
    "        'valid-table',\n",
    "        value='Energy.MLDataValid'),\n",
    "    test_table = dsl.PipelineParam(\n",
    "        'test-table',\n",
    "        value='Energy.MLDataTest'),\n",
    "    train_file = dsl.PipelineParam(\n",
    "        'train-file',\n",
    "        value='gs://energyforecast/data/csv/MLDataTrain.csv'),\n",
    "    valid_file = dsl.PipelineParam(\n",
    "        'valid-file',\n",
    "        value='gs://energyforecast/data/csv/MLDataValid.csv'),\n",
    "    test_file = dsl.PipelineParam(\n",
    "        'test-file',\n",
    "        value='gs://energyforecast/data/csv/MLDataTest.csv')):\n",
    "    \n",
    "    CopTableOp = kfp.components.load_component('copy-table.component.yaml')\n",
    "    ExpTableOp = kfp.components.load_component('export-table.component.yaml')\n",
    "    DataPrepOp = kfp.components.load_component('run-git-python-script.component.yaml')\n",
    "    TrainModelOp = kfp.components.load_component('train-git-cmle-model.component.yaml')\n",
    "    \n",
    "    \n",
    "    ct_op = CopTableOp(\n",
    "        dataset).apply(gcp.use_gcp_secret('user-gcp-sa'))\n",
    "    dp_op = DataPrepOp(\n",
    "        ct_op.output,\n",
    "        code_repo,\n",
    "        code_folder,\n",
    "        data_prep_script,\n",
    "        data_prep_args).apply(gcp.use_gcp_secret('user-gcp-sa'))\n",
    "    tr_et_op = ExpTableOp(\n",
    "        dp_op.output,\n",
    "        train_table,\n",
    "        train_file).apply(gcp.use_gcp_secret('user-gcp-sa'))\n",
    "    va_et_op = ExpTableOp(\n",
    "        dp_op.output,\n",
    "        valid_table,\n",
    "        valid_file).apply(gcp.use_gcp_secret('user-gcp-sa'))\n",
    "    te_et_op = ExpTableOp(\n",
    "        dp_op.output,\n",
    "        test_table,\n",
    "        test_file).apply(gcp.use_gcp_secret('user-gcp-sa'))\n",
    "    tm_op = TrainModelOp(\n",
    "        tr_et_op.output,\n",
    "        va_et_op.output,\n",
    "        code_repo,\n",
    "        code_folder,\n",
    "        project,\n",
    "        bucket,\n",
    "        package_folder,\n",
    "        cmle_folder,\n",
    "        scale_tier,\n",
    "        python_module,\n",
    "        region,\n",
    "        runtime_version,\n",
    "        cmle_args).apply(gcp.use_gcp_secret('user-gcp-sa'))\n",
    "    \n",
    "compiler.Compiler().compile(basic_bq_pipeline, 'energy-forecasting.tar.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
